import datetime
import backtrader as bt
from backtrader_plotting import Bokeh
import pandas as pd
import akshare as ak


def get_data(code):
    df = ak.stock_zh_a_daily(symbol=code, adjust='qfq')
    # df.date = pd.to_datetime(df.date, format='%Y-%m-%d', utc=True)
    # df.set_index('date', inplace=True)
    df.index = pd.to_datetime(df.date)
    df.sort_index(inplace=True)

    return df


class TestStrategy(bt.Strategy):
    params = (
        ('buydate', 21),
        ('holdtime', 6),
    )

    def next(self):
        if len(self.data) == self.p.buydate:
            self.buy(self.datas[0], size=None)

        if len(self.data) == self.p.buydate + self.p.holdtime:
            self.sell(self.datas[0], size=None)


if __name__ == '__main__':
    cerebro = bt.Cerebro()

    cerebro.addstrategy(TestStrategy, buydate=3)

    codes = ['sh600276', 'sh600519', 'sh603288']
    # 恒瑞医药
    data = get_data(codes[0])
    st_date = datetime.datetime(2023, 3, 1)
    end_date = datetime.datetime(2023, 6, 21)
    datafeed = bt.feeds.PandasData(dataname=data, fromdate=st_date, todate=end_date)
    cerebro.adddata(datafeed, name='600276.SH')

    cerebro.run()

    b = Bokeh(style='bar', plot_mode='single')
    cerebro.plot(b)
